import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import cosine_similarity

# 假设有多个批次
num_batches = 10  # 批次数量
batch_size = 100  # 每个批次的大小
feature_size = 512  # 特征维度

# 随机生成多个批次的数据，实际使用时替换为真实特征
A_batches = [np.random.randn(batch_size, feature_size) for _ in range(num_batches)]  # A 网络的特征
B_batches = [np.random.randn(batch_size, feature_size) for _ in range(num_batches)]  # B 网络的特征

# 计算每个批次的余弦相似度，并计算每个批次的平均余弦相似度
avg_cos_sim = []
for i in range(num_batches):
    A = A_batches[i]
    B = B_batches[i]
    # 计算余弦相似度矩阵 [batch_size, batch_size]
    cos_sim = cosine_similarity(A, B)
    # 计算每个批次的平均余弦相似度
    avg_cos_sim.append(np.mean(cos_sim))

# 计算所有批次的平均余弦相似度
final_avg_cos_sim = np.mean(avg_cos_sim)

# 绘制条形图
plt.figure(figsize=(10, 6))
plt.bar(range(1, num_batches + 1), avg_cos_sim, color='skyblue', label='Batch Average Cosine Similarity')
plt.axhline(y=final_avg_cos_sim, color='red', linestyle='--', label=f'Final Average: {final_avg_cos_sim:.4f}')
plt.xlabel('Batch Number')
plt.ylabel('Average Cosine Similarity')
plt.title('Average Cosine Similarity Between A and B Features Across Batches')
plt.xticks(range(1, num_batches + 1))  # 设置x轴标签为批次号
plt.legend()  # 显示图例
plt.show()